# 1-定义客户端
from langchain_core.prompts import ChatMessagePromptTemplate
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from pydantic import SecretStr, BaseModel, Field

from settings import DASHSCOPE_API_KEY

# 1-定义客户端
llm = ChatOpenAI(
    model="qwen-max",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    api_key=SecretStr(DASHSCOPE_API_KEY),
    streaming=True
)
llm_no_stream = ChatOpenAI(
    model="qwen-max",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    api_key=SecretStr(DASHSCOPE_API_KEY),
    streaming=False  # 先关闭流式，确保返回 AIMessage
)

# 2-定义对话上下文
system_message_template = ChatMessagePromptTemplate.from_template(
    template="你是一位{role}专家，擅长回答{domain}领域的问题",
    role="system"
)

human_message_template = ChatMessagePromptTemplate.from_template(
    template="用户问题：{question}",
    role="user"
)

llm_messages = [
    system_message_template,
    human_message_template,
]

# 3-定义日志的打印函数
def llm_stream_print_result(llm, chat_prompt_detail):
    response = llm.stream(chat_prompt_detail)
    print("-----Hold on, LLM 正在回答！-----")
    for chunk in response:
        # 打印不换行
        print(chunk.content, end="")

# 4-定义工具
## 4.1 定义参数对应的类描述
class InputArgus(BaseModel):
    a: int = Field(description="数字1")
    b: int = Field(description="数字2")


## 4.2 定义工具方法
@tool(description="计算两个数字的和",
      args_schema=InputArgus)
def add(a: int, b: int) -> int:
    return a + b

## 4.3 定义工具集
def calculate_calc_tool():
    return [add]